import copy
import numpy as np
import os
from dataflow.my_io import read_points
from wwengine.utils import bbox3d2bevcorners, box_collision_test, remove_pts_in_bboxes, limit_period, Priority
from ..builder import TRANSFORMS
from .transfrom_base import TransfromBase

class BaseSampler():
    def __init__(self, sampled_list, shuffle=True):
        self.total_num = len(sampled_list)
        self.sampled_list = np.array(sampled_list)
        self.indices = np.arange(self.total_num)
        if shuffle:
            np.random.shuffle(self.indices)
        self.shuffle = shuffle
        self.idx = 0

    def sample(self, num):
        if self.idx + num < self.total_num:
            ret = self.sampled_list[self.indices[self.idx:self.idx + num]]
            self.idx += num
        else:
            ret = self.sampled_list[self.indices[self.idx:]]
            self.idx = 0
            if self.shuffle:
                np.random.shuffle(self.indices)
        return ret


@TRANSFORMS.register_module()
class Dbsample(TransfromBase):
    def __init__(self, sample_minpts, sample_groups, class2label, data_root, db_infos):
        super(Dbsample, self).__init__(['DET'],Priority.HIGHEST)
        self.class2label = class2label
        self.data_root = data_root
        # 如果数据增强配置中存在sample配置才进行sample
        db_sampler = {}
        sample_groups_exist = {}
        db_infos = self.filter_db(db_infos, filter_thrs=sample_minpts)  # 把点数少的sample滤除
        for cat_name in self.class2label:  # 待识别类型,且sample中有的,才进行数据增强
            if cat_name in db_infos:
                db_sampler[cat_name] = BaseSampler(db_infos[cat_name], shuffle=True)
                sample_groups_exist[cat_name] = sample_groups[cat_name]  # samples中已有的数据,才能进行采样
        self.db_sampler = db_sampler
        self.sample_groups = sample_groups_exist
        # self.tasks = ['DET']  # 当前aug支持的任务
        # self.priority = Priority.HIGHEST  # 当前aug的优先级

    def filter_db(self, db_infos, filter_thrs=None):
        # 1. filter_by_min_points, dict(Car=5, Pedestrian=10, Cyclist=10)
        filtered_db_infos = dict()
        for cat in self.class2label:
            if cat in db_infos:
                filter_thr = filter_thrs[cat]
                filtered_db_infos[cat] = [item for item in db_infos[cat] if item['num_points_in_gt'] >= filter_thr]

        return filtered_db_infos

    def forward(self, data_dict):
        lidar_pts, gt_bboxes_3d = data_dict['lidar_pts'], data_dict['gt_bboxes_3d']
        gt_labels, gt_names = data_dict['gt_labels'], data_dict['gt_names']

        sampled_pts, sampled_names, sampled_labels = [], [], []
        sampled_bboxes = []

        avoid_coll_boxes = copy.deepcopy(gt_bboxes_3d)
        for name, v in self.sample_groups.items():
            # 1. calculate sample numbers
            sampled_num = v - np.sum(gt_names == name)
            if sampled_num <= 0:
                continue
            # 2. sample databases bboxes
            sampled_cls_list = self.db_sampler[name].sample(sampled_num)  # 从样本库中取出sampled_num个样本信息
            sampled_cls_bboxes = np.array([item['box3d_lidar'] for item in sampled_cls_list],
                                          dtype=np.float32)  # 从样本信息中提取box

            # 3. box_collision_test
            # box转bev视角的2dbox的顶点
            avoid_coll_boxes_bv_corners = bbox3d2bevcorners(avoid_coll_boxes)  # ori boxes
            sampled_cls_bboxes_bv_corners = bbox3d2bevcorners(sampled_cls_bboxes)  # 来自样本库的new boxes
            coll_query_matrix = np.concatenate([avoid_coll_boxes_bv_corners, sampled_cls_bboxes_bv_corners],
                                               axis=0)  # ori+new
            coll_mat = box_collision_test(coll_query_matrix, coll_query_matrix)  # 碰撞检测矩阵
            n_gt, tmp_bboxes = len(avoid_coll_boxes_bv_corners), []
            for i in range(n_gt, len(coll_mat)):
                if any(coll_mat[i]):  # any()函数用于判断可迭代参数iterable是否全部为False，返回False；如果有一个为True，则返回True
                    coll_mat[i] = False
                    coll_mat[:, i] = False
                else:  # 如果新增的和之前的不碰撞
                    cur_sample = sampled_cls_list[i - n_gt]
                    pt_path = os.path.join(self.data_root, cur_sample['path'])
                    sampled_pts_cur = read_points(pt_path, dim=3, datatype=np.float32)
                    sampled_pts_cur[:, :3] += cur_sample['box3d_lidar'][:3]  # sample保存是以原点为中心,所以要加上偏移量
                    sampled_pts.append(sampled_pts_cur)
                    sampled_names.append(cur_sample['name'])
                    sampled_labels.append(self.class2label[cur_sample['name']])
                    sampled_bboxes.append(cur_sample['box3d_lidar'])
                    tmp_bboxes.append(cur_sample['box3d_lidar'])
            if len(tmp_bboxes) == 0:
                tmp_bboxes = np.array(tmp_bboxes).reshape(-1, 7)
            else:
                tmp_bboxes = np.array(tmp_bboxes)
            avoid_coll_boxes = np.concatenate([avoid_coll_boxes, tmp_bboxes], axis=0)

        # merge sampled database
        # remove raw points in sampled_bboxes firstly
        if len(sampled_bboxes) > 0:  # 如果找到sample，就把sample box内的原点云去掉，换成sample的点云
            # # # 测试:新增一个box測試是否box内的点被去掉
            # sampled_bboxes.append(np.array([0,0,0,20,5,5,0.75]))
            # avoid_coll_boxes = np.concatenate([avoid_coll_boxes, np.array([[0,0,0,20,5,5,0.75]])], axis=0)
            # gt_labels = np.concatenate([gt_labels, np.array([2])], axis=0)
            # gt_names = np.concatenate([gt_names, np.array(['Car'])], axis=0)

            lidar_pts = remove_pts_in_bboxes(lidar_pts, np.stack(sampled_bboxes, axis=0))
            # lidar_pts = np.concatenate([lidar_pts, np.concatenate(sampled_pts, axis=0)], axis=0)
            lidar_pts = np.concatenate([np.concatenate(sampled_pts, axis=0), lidar_pts], axis=0)
            gt_bboxes_3d = avoid_coll_boxes.astype(np.float32)
            gt_labels = np.concatenate([gt_labels, np.array(sampled_labels)], axis=0)
            gt_names = np.concatenate([gt_names, np.array(sampled_names)], axis=0)

        # 更新data_dict
        data_dict.update({'lidar_pts': lidar_pts})
        data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
        data_dict.update({'gt_labels': gt_labels})
        data_dict.update({'gt_names': gt_names})
        data_dict['old_bboxes_size'] = len(data_dict['gt_bboxes_3d'])

        return data_dict


@TRANSFORMS.register_module()
class ObjectNoise(TransfromBase):
    def __init__(self, num_try, translation_std, rot_range):
        super(ObjectNoise, self).__init__(['DET'], Priority.NORMAL)
        self.num_try = num_try
        self.translation_std = translation_std
        self.rot_range = rot_range

    def forward(self, data_dict):
        lidar_pts, gt_bboxes_3d = data_dict['lidar_pts'], data_dict['gt_bboxes_3d']
        n_bbox = len(gt_bboxes_3d)

        # 1. generate rotation vectors and rotation matrices
        trans_vec = np.random.normal(scale=self.translation_std, size=(n_bbox, self.num_try, 3)).astype(np.float32)
        rot_angle = np.random.uniform(self.rot_range[0], self.rot_range[1], size=(n_bbox, self.num_try)).astype(np.float32)
        rot_cos, rot_sin = np.cos(rot_angle), np.sin(rot_angle)
        # in fact, - rot_angle
        rot_mat = np.array([[rot_cos, -rot_sin],
                            [rot_sin, rot_cos]])  # (2, 2, n_bbox, num_try)
        rot_mat = np.transpose(rot_mat, (2, 3, 1, 0))  # (n_bbox, num_try, 2, 2)

        # 2. generate noise for each bbox and the points inside the bbox.
        bev_corners = bbox3d2bevcorners(gt_bboxes_3d)  # (n_bbox, 4, 2) # for collision test
        masks = remove_pts_in_bboxes(lidar_pts, gt_bboxes_3d, rm=False)  # identify which point should be added noise
        gt_bboxes_3d, lidar_pts = object_noise_core(pts=lidar_pts,
                                                    gt_bboxes_3d=gt_bboxes_3d,
                                                    bev_corners=bev_corners,
                                                    trans_vec=trans_vec,
                                                    rot_angle=rot_angle,
                                                    rot_mat=rot_mat,
                                                    masks=masks)
        data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
        data_dict.update({'lidar_pts': lidar_pts})

        return data_dict


# @numba.jit(nopython=True)
def object_noise_core(pts, gt_bboxes_3d, bev_corners, trans_vec, rot_angle, rot_mat, masks):
    '''
    pts: (N, 4)
    gt_bboxes_3d: (n_bbox, 7)
    bev_corners: ((n_bbox, 4, 2))
    trans_vec: (n_bbox, num_try, 3)
    rot_mat: (n_bbox, num_try, 2, 2)
    masks: (N, n_bbox), bool
    return: gt_bboxes_3d, pts
    '''
    # 1. select the noise of num_try for each bbox under the collision test
    n_bbox, num_try = trans_vec.shape[:2]

    # succ_mask: (n_bbox, ), whether each bbox can be added noise successfully. -1 denotes failure.
    succ_mask = -np.ones((n_bbox,), dtype=np.int_)
    for i in range(n_bbox):
        for j in range(num_try):
            cur_bbox = bev_corners[i] - np.expand_dims(gt_bboxes_3d[i, :2], 0)  # (4, 2) - (1, 2) -> (4, 2)
            rot = np.zeros((2, 2), dtype=np.float32)
            rot[:] = rot_mat[i, j]  # (2, 2)
            trans = trans_vec[i, j]  # (3, )
            cur_bbox = cur_bbox @ rot
            cur_bbox += gt_bboxes_3d[i, :2]
            cur_bbox += np.expand_dims(trans[:2], 0)  # (4, 2)
            coll_mat = box_collision_test(np.expand_dims(cur_bbox, 0), bev_corners)
            coll_mat[0, i] = False
            if coll_mat.any():
                continue
            else:
                bev_corners[i] = cur_bbox  # update the bev_corners when adding noise succseefully.
                succ_mask[i] = j
                break
    # 2. points and bboxes noise
    visit = {}
    for i in range(n_bbox):
        jj = succ_mask[i]
        if jj == -1:
            continue
        cur_trans, cur_angle = trans_vec[i, jj], rot_angle[i, jj]
        cur_rot_mat = np.zeros((2, 2), dtype=np.float32)
        cur_rot_mat[:] = rot_mat[i, jj]
        for k in range(len(pts)):
            if masks[k][i] and k not in visit:
                cur_pt = pts[k]  # (4, )
                cur_pt_xyz = np.zeros((1, 3), dtype=np.float32)
                cur_pt_xyz[0] = cur_pt[:3] - gt_bboxes_3d[i][:3]
                tmp_cur_pt_xy = np.zeros((1, 2), dtype=np.float32)
                tmp_cur_pt_xy[:] = cur_pt_xyz[:, :2]
                cur_pt_xyz[:, :2] = tmp_cur_pt_xy @ cur_rot_mat  # (1, 2)
                cur_pt_xyz[0] = cur_pt_xyz[0] + gt_bboxes_3d[i][:3]
                cur_pt_xyz[0] = cur_pt_xyz[0] + cur_trans[:3]
                cur_pt[:3] = cur_pt_xyz[0]
                visit[k] = 1

        gt_bboxes_3d[i, :3] += cur_trans[:3]
        gt_bboxes_3d[i, 6] += cur_angle

    return gt_bboxes_3d, pts


@TRANSFORMS.register_module()
class RandomFlip(TransfromBase):
    def __init__(self, random_flip_ratio=0.5):
        super(RandomFlip, self).__init__(['DET'], Priority.NORMAL)
        self.random_flip_ratio = random_flip_ratio

    def forward(self, data_dict):
        random_flip_state = np.random.choice([True, False], p=[self.random_flip_ratio, 1 - self.random_flip_ratio])
        if random_flip_state:
            lidar_pts = data_dict['lidar_pts']
            lidar_pts[:, 1] = -lidar_pts[:, 1]
            gt_bboxes_3d = data_dict['gt_bboxes_3d']
            gt_bboxes_3d[:, 1] = -gt_bboxes_3d[:, 1]
            gt_bboxes_3d[:, 6] = -gt_bboxes_3d[:, 6] + np.pi
            data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
            data_dict.update({'lidar_pts': lidar_pts})
        return data_dict


@TRANSFORMS.register_module()
class GlobalRotScaleTrans(TransfromBase):
    def __init__(self, rot_range, scale_ratio_range, translation_std):
        super(GlobalRotScaleTrans, self).__init__(['DET'], Priority.NORMAL)
        self.rot_range = rot_range
        self.scale_ratio_range = scale_ratio_range
        self.translation_std = translation_std

    def forward(self, data_dict):
        lidar_pts, gt_bboxes_3d = data_dict['lidar_pts'], data_dict['gt_bboxes_3d']

        # 1. rotation
        rot_angle = np.random.uniform(self.rot_range[0], self.rot_range[1])
        rot_cos, rot_sin = np.cos(rot_angle), np.sin(rot_angle)
        # in fact,  rot_angle
        rot_mat = np.array([[rot_cos, -rot_sin],
                            [rot_sin, rot_cos]])  # (2, 2)
        # # 1.1 bbox rotation   x1 = R @ x0 ->  x1.T = x0_T @ R.T
        gt_bboxes_3d[:, :2] = gt_bboxes_3d[:, :2] @ rot_mat.T
        gt_bboxes_3d[:, 6] += rot_angle
        # 1.2 point rotation
        lidar_pts[:, :2] = lidar_pts[:, :2] @ rot_mat.T

        # 2. scaling
        scale_fator = np.random.uniform(self.scale_ratio_range[0], self.scale_ratio_range[1])
        gt_bboxes_3d[:, :6] *= scale_fator
        lidar_pts[:, :3] *= scale_fator

        # 3. translation
        trans_factor = np.random.normal(scale=self.translation_std, size=(1, 3))
        gt_bboxes_3d[:, :3] += trans_factor
        lidar_pts[:, :3] += trans_factor
        data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
        data_dict.update({'lidar_pts': lidar_pts})
        return data_dict


@TRANSFORMS.register_module()
class RangeFilter(TransfromBase):
    def __init__(self, range):
        super(RangeFilter, self).__init__(['DET'], Priority.LOWEST)
        self.point_range = range
        assert range[3]-range[0] > 4.0 and range[4]-range[1] > 4.0
        self.object_range = range.copy()
        self.object_range[0] = self.object_range[0] + 2.0
        self.object_range[1] = self.object_range[1] + 2.0
        self.object_range[3] = self.object_range[3] - 2.0
        self.object_range[4] = self.object_range[4] - 2.0

    def forward(self, data_dict):
        # point_range_filter
        pts = data_dict['lidar_pts']
        flag_x_low = pts[:, 0] > self.point_range[0]
        flag_y_low = pts[:, 1] > self.point_range[1]
        flag_z_low = pts[:, 2] > self.point_range[2]
        flag_x_high = pts[:, 0] < self.point_range[3]
        flag_y_high = pts[:, 1] < self.point_range[4]
        flag_z_high = pts[:, 2] < self.point_range[5]
        keep_mask = flag_x_low & flag_y_low & flag_z_low & flag_x_high & flag_y_high & flag_z_high
        pts = pts[keep_mask]
        data_dict.update({'lidar_pts': pts})

        gt_bboxes_3d, gt_labels = data_dict['gt_bboxes_3d'], data_dict['gt_labels']
        gt_names = data_dict['gt_names']

        # object_range_bev_filter
        flag_x_low = gt_bboxes_3d[:, 0] > self.object_range[0]
        flag_y_low = gt_bboxes_3d[:, 1] > self.object_range[1]
        flag_x_high = gt_bboxes_3d[:, 0] < self.object_range[3]
        flag_y_high = gt_bboxes_3d[:, 1] < self.object_range[4]
        keep_mask = flag_x_low & flag_y_low & flag_x_high & flag_y_high

        gt_bboxes_3d, gt_labels = gt_bboxes_3d[keep_mask], gt_labels[keep_mask]
        gt_names = gt_names[keep_mask]
        gt_bboxes_3d[:, 6] = limit_period(gt_bboxes_3d[:, 6], 0.5, 2 * np.pi)
        data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
        data_dict.update({'gt_labels': gt_labels})
        data_dict.update({'gt_names': gt_names})

        return data_dict


@TRANSFORMS.register_module()
class PointsShuffle(TransfromBase):
    def __init__(self, random_flip_ratio=0.5):
        super(PointsShuffle, self).__init__(['DET'], Priority.NORMAL)

    def forward(self, data_dict):
        pts = data_dict['lidar_pts']
        indices = np.arange(0, len(pts))
        np.random.shuffle(indices)
        pts = pts[indices]
        data_dict.update({'lidar_pts': pts})
        return data_dict



# def filter_bboxes_with_labels(data_dict, label=-1):
#     '''
#     data_dict: dict(pts, gt_bboxes_3d, gt_labels, gt_names, difficulty)
#     label: int
#     '''
#     gt_bboxes_3d, gt_labels = data_dict['gt_bboxes_3d'], data_dict['gt_labels']
#     gt_names, difficulty = data_dict['gt_names'], data_dict['difficulty']
#     idx = gt_labels != label
#     gt_bboxes_3d = gt_bboxes_3d[idx]
#     gt_labels = gt_labels[idx]
#     gt_names = gt_names[idx]
#     difficulty = difficulty[idx]
#     data_dict.update({'gt_bboxes_3d': gt_bboxes_3d})
#     data_dict.update({'gt_labels': gt_labels})
#     data_dict.update({'gt_names': gt_names})
#     data_dict.update({'difficulty': difficulty})
#     return data_dict

